# -*- coding: utf-8 -*-  
'''
LDA模型

Created on 2021年8月21日
@author: luoyi
'''
import tensorflow as tf

import utils.conf as conf
from utils.iexicon import LiteWordsWarehouse
from models.abstract_nets import AbstractModel
from models.lda.parts import GibbsLayer
from models.lda.losses import LDALoss
from models.lda.metrics import LDAMetric
from models.lda.callbacks import AutoSaveLDACallable


#    LDA模型
class LDA(AbstractModel):
    '''LDA模型
    
    '''
    def __init__(self,
                 name='LDA',
                 K=conf.LDA.get_k(),
                 W=LiteWordsWarehouse.instance().words_count(),
                 max_doc_words=conf.LDA.get_max_doc_words(),
                 input_shape=(None, conf.LDA.get_max_doc_words())):
        '''
            @param K: 主题数
            @param max_doc_words: 每篇文档最大词数
        '''
        self._K = K
        self._W = W
        self._max_doc_words = max_doc_words
        
        super(LDA, self).__init__(learning_rate=0.0001, name=name, auto_assembling=True, is_build=True, input_shape=input_shape)
        pass
    
    def append_callbacks(self):
        return [AutoSaveLDACallable(n_batch=conf.LDA.get_auto_save_batch(), 
                                    m_epoch=conf.LDA.get_auto_save_epoch(), 
                                    wn_topic_word_path=conf.LDA.get_wn_topic_word_path(), 
                                    prob_topic_word_path=conf.LDA.get_prob_topic_word_path())]
    
    def create_optimizer(self, learning_rate=0.001):
        return tf.optimizers.Adam(learning_rate)
    
    def create_loss(self):
        return LDALoss()
    
    def create_metrics(self):
        return [LDAMetric(prob_layer=self._prob_layer, 
                          K=self._K,
                          max_doc_words=self._max_doc_words)]
    
    #    装配网络
    def assembling(self):
        self._prob_layer = GibbsLayer(K=self._K,
                                      W=self._W,
                                      max_doc_words=self._max_doc_words)
        pass
    
    def call(self, inputs, training=None, mask=None):
        out = self._prob_layer(inputs)
        return out
    
    pass